import tensorflow.compat.v1 as tf
import tensorflow.examples.tutorials.mnist.input_data as input_data

"""
MNIST 手写数字识别
训练集: 55000
验证集: 5000
测试集: 10000

one hot: 独热编码, 一种稀疏向量, 其中一个元素设为1, 其余元素全设为0
    np.argmax(mnist.train.labels[index])
"""
mnist = input_data.read_data_sets("MNIST_data/", one_hot=True)

'''模型构建'''
# 定义占位符
# minst每张图片共有28*28=784个像素点
x = tf.placeholder(tf.float32, [None, 784], name="X")
# 0-9 一共十个数字 => 十个类别
y = tf.placeholder(tf.float32, [None, 10], name="Y")

# 定义变量
w = tf.Variable(tf.random_normal([784, 10]), name="W")
b = tf.Variable(tf.zeros([10]), name="B")

# 用单个神经元构建神经网络
# 向前计算
forward = tf.matmul(x, w) + b
# 当我们处理多酚类任务时, 通常需要使用Softmax Regression模型
# 会对每一类别分别估算出一个概率
pred = tf.nn.softmax(forward)

'''训练模型'''
# 设置训练参数
# 训练轮数
train_epoch = 50
# 单词训练样本数(批次大小)
batch_size = 100
#  一轮训练有多少批次
total_batch = int(mnist.train.num_examples / batch_size)
# 显示粒度
display_step = 1
# 学习率
learning_rate = 0.01

# 定义损失函数
# 定义交叉熵损失函数
loss_function = tf.reduce_mean(
    -tf.reduce_sum(y * tf.log(pred), reduction_indices=1))

# 选择优化器
# 梯度下降优化器
optimizer = tf.train.GradientDescentOptimizer(learning_rate).minimize(loss_function)

# 定义精准率
# 检查预测类别tf.argmax(pred, 1)与实际类别tf.argmax(y, 1)的匹配情况
correct_prediction = tf.equal(tf.argmax(pred, 1), tf.argmax(y, 1))
# 准确率, 将布尔值装华为浮点数, 并计算平均值
accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

# 声明会话
sess = tf.Session()
# 变量初始化
init = tf.global_variables_initializer()
sess.run(init)

'''开始训练'''
for epoch in range(train_epoch):
    for batch in range(total_batch):
        # 读取批次数据
        xs, ys = mnist.train.next_batch(batch_size)
        # 执行批次训练
        sess.run(optimizer, feed_dict={x: xs, y: ys})

    # total_batch个批次训练完成后, 使用验证数据计算误差与准确率, 验证集没有分批
    loss, acc = sess.run([loss_function, accuracy],
                         feed_dict={x: mnist.validation.images, y: mnist.validation.labels})
    # 打印训练过程中的详情信息
    if (epoch + 1) % display_step == 0:
        print("Train Epoch:", "%02d" % (epoch + 1),
              "Loss=", "{:.9f}".format(loss),
              "Accuracy=", "{:.4f}".format(acc))

print("Train Finished!")
